"""
Author: Morphlng
Date: 2024-04-25 10:29:42
LastEditTime: 2024-04-29 21:44:11
LastEditors: Morphlng
Description: Low-level walker control for Carla simulator
FilePath: /DrivingGym/src/driving_gym/environment/agent/action/carla_walker_action.py
"""

from __future__ import annotations

import carla
import gymnasium as gym
import numpy as np

from driving_gym.environment.agent.action.action_interface import (
    AbstractAction,
    ActionInterface,
)
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class DirectAction(AbstractAction):
    def __init__(self, adapter, actor_id, action, duration=1):
        super().__init__(action, duration)
        self.adapter = adapter
        self.actor_id = actor_id

    @override(AbstractAction)
    def run_step(self):
        """Return a carla control signal dict."""
        self.duration -= 1

        transform = self.adapter.get_transform(self.actor_id)
        yaw = transform.rotation.yaw

        direction = self.action["direction"]
        x, y = self.action_transform(yaw, self.get_direction_vector(direction))

        return {
            "type": "walker",
            "speed": self.action["speed"],
            "direction": carla.Vector3D(x, y, 0),
        }

    @staticmethod
    def action_transform(yaw, direction):
        yaw = np.radians(-yaw)
        rotation_matrix = np.array(
            [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]]
        )
        inverse_rotation_matrix = np.transpose(rotation_matrix)
        return np.dot(inverse_rotation_matrix, direction)

    @staticmethod
    def get_direction_vector(direction_str):
        """Convert direction string to carla.Vector3D"""

        direction_map = {
            "stay": (0, 0),
            "front": (1, 0),
            "left": (0, -1),
            "right": (0, 1),
        }

        x, y = direction_map[direction_str]
        return np.array([x, y])


class CarlaWalkerAction(ActionInterface):
    def __init__(self, action_config: dict, adapter: AdapterInterface):
        """Initialize the action converter for walker action space

        Args:
            action_config (dict): A dictionary of action config
        """
        super().__init__(action_config, adapter)

        # Direction is fixed to stay, front, left, right
        self._direction_set = ["stay", "front", "left", "right"]
        # Speed is adjustable
        self._speed_set = self.action_config.get("speed_set", [0, 1, 2, 3])

    @override(ActionInterface)
    def convert_single_action(self, action):
        """Convert the action of a model output to an AbstractAction instance

        Args:
            action: Action input from the model

        Returns:
            DirectAction: A direct action instance
        """
        action_dict = {
            "direction": self._direction_set[action[0]],
            "speed": self._speed_set[action[1]],
        }
        return DirectAction(self.adapter, self.actor_id, action_dict)

    @override(ActionInterface)
    def get_action_space(self):
        return gym.spaces.MultiDiscrete(
            [len(self._direction_set), len(self._speed_set)]
        )

    @override(ActionInterface)
    def get_mask_space(self):
        return None
